import factory
from app.models.knowledge_base import KnowledgeBase, Document
from app.models.system_prompt import SystemPromptTemplate
from app import db

class KnowledgeBaseFactory(factory.alchemy.SQLAlchemyModelFactory):
    class Meta:
        model = KnowledgeBase
        sqlalchemy_session = db.session

    name = factory.Sequence(lambda n: f'Test KB {n}')
    description = factory.Faker('text')
    embedding_model = 'test-model'

class DocumentFactory(factory.alchemy.SQLAlchemyModelFactory):
    class Meta:
        model = Document
        sqlalchemy_session = db.session

    title = factory.Faker('sentence')
    content = factory.Faker('text')
    knowledge_base = factory.SubFactory(KnowledgeBaseFactory) 